# -*- coding: utf-8 -*-
"""
author:LTH
data:
"""
import numpy as np
import torch
from torch import nn
from torch.backends import cudnn
from torch.utils.data import DataLoader
from tqdm import tqdm

from config import GetConfig
from datalist import train_dataset
from model import EModel, ArcMarginProduct,ResNet,ArcMarginModel


class ArcFace(object):
    def __init__(self):
        self.args = GetConfig()
        print(f"-----------{self.args.project_name}-------------")

        use_cuda = True
        self.device = torch.device("cuda" if use_cuda else "cpu")

        kwargs = {"num_workers": 0, "pin_memory": True} if use_cuda else {"num_workers": 0, "pin_memory": False}
        '''
        构造DataLoader
        '''
        self.train_dataset = train_dataset
        self.train_dataloader = DataLoader(self.train_dataset, batch_size=self.args.train_batch_size, shuffle=True,
                                           **kwargs)

        self.test_dataloader = DataLoader(self.train_dataset, batch_size=2, shuffle=False,
                                           **kwargs)
        '''
        定义模型
        '''
        self.beddingNet = ResNet().to(self.device)
        self.head = ArcMarginModel().to(self.device)


        '''
        CUDA加速
        '''
        if use_cuda:
            self.beddingNet = torch.nn.DataParallel(self.beddingNet, device_ids=range(torch.cuda.device_count()))
            self.head = torch.nn.DataParallel(self.head, device_ids=range(torch.cuda.device_count()))
            cudnn.enabled = True
            cudnn.benchmark = True

        self.beddingNet.load_state_dict(torch.load("bedding_best.pth")["model_state_dict"])
        self.head.load_state_dict(torch.load("head_best.pth")['model_state_dict'])

        '''
        构造loss目标函数
        选择优化器
        学习率变化选择
        '''
        self.optimizer = torch.optim.Adam([
            {'params': self.beddingNet.parameters(), 'lr': self.args.lr},
            {'params': self.head.parameters(), 'lr': self.args.lr},
        ])
        self.criterion = nn.CrossEntropyLoss()
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=5, eta_min=1e-5)

        # try:
        for epoch in range(1, self.args.epochs + 1):
            # self.train(epoch)
            # torch.save({"model_state_dict":self.beddingNet.state_dict()},"bedding_best.pth")
            # torch.save({"model_state_dict": self.head.state_dict()}, "head_best.pth")
            self.test(epoch)
        # except Exception as e:
        #     torch.cuda.empty_cache()
        #     print("stop model training")

    def train(self, epoch):
        self.beddingNet.train()
        self.head.train()

        average_loss = []
        total_correct = 0
        total_num = 0
        pbar = tqdm(self.train_dataloader, desc=f'Train Epoch {epoch}/{self.args.epochs}')
        for image, target in pbar:
            self.optimizer.zero_grad()
            image, target = image.to(self.device), target.to(self.device)
            feature = self.beddingNet(image)
            out= self.head(feature, target)
            loss=self.criterion(out,target)
            loss.backward()
            self.optimizer.step()
            average_loss.append(loss.item())
            _, pred = out.max(1)
            num_correct = (pred == target).sum().item()
            total_correct += num_correct
            total_num += image.shape[0]
            train_acc = total_correct / total_num
            pbar.set_description(f"Epoch {epoch}"
                                 f"\t train_acc:{train_acc}"
                                 f"\t loss: {np.mean(average_loss)}")

    def test(self,epoch):
        self.beddingNet.eval()
        self.head.eval()
        pbar=tqdm(self.test_dataloader, desc=f'Train Epoch {epoch}/{self.args.epochs}')
        for image,target in pbar:
            image, target = image.to(self.device), target.to(self.device)
            feature = self.beddingNet(image)
            print(cal_cos(feature[0],feature[1]))
            print(target[0],target[1])


def cal_cos(feature1,feature2):
    return (torch.sum(feature2.T*feature1))/(torch.norm(feature2,2)*torch.norm(feature1,2))
model = ArcFace()
